# -*- coding: utf-8 -*-

import subprocess
import math
import torch
import torch.nn as nn
from torch.autograd import Variable
import numpy as np

import float_tensor_pb2


def conv_bn(inp, oup, stride):
    return nn.Sequential(
        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
        nn.BatchNorm2d(oup),
        nn.ReLU(inplace=True)
    )


def conv_1x1_bn(inp, oup):
    return nn.Sequential(
        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
        nn.BatchNorm2d(oup),
        nn.ReLU(inplace=True)
    )


class InvertedResidual(nn.Module):
    def __init__(self, inp, oup, stride, expand_ratio):
        super(InvertedResidual, self).__init__()
        self.stride = stride
        assert stride in [1, 2]

        self.use_res_connect = self.stride == 1 and inp == oup

        self.conv = nn.Sequential(
            # pw
            nn.Conv2d(inp, inp * expand_ratio, 1, 1, 0, bias=False),
            nn.BatchNorm2d(inp * expand_ratio),
            nn.ReLU(inplace=True),
            # dw
            nn.Conv2d(inp * expand_ratio, inp * expand_ratio, 3, stride, 1, groups=inp * expand_ratio, bias=False),
            nn.BatchNorm2d(inp * expand_ratio),
            nn.ReLU(inplace=True),
            # pw-linear
            nn.Conv2d(inp * expand_ratio, oup, 1, 1, 0, bias=False),
            nn.BatchNorm2d(oup),
        )

    def forward(self, x):
        if self.use_res_connect:
            return x + self.conv(x)
        else:
            return self.conv(x)


class MobileNetV2(nn.Module):
    def __init__(self, n_class=1000, input_size=224, width_mult=1.):
        super(MobileNetV2, self).__init__()
        # setting of inverted residual blocks
        self.inverted_residual_setting = [
            # t, c, n, s
            [1, 16, 1, 1],
            [6, 24, 2, 2],
            [6, 32, 3, 2],
            [6, 64, 4, 2],
            [6, 96, 3, 1],
            [6, 160, 3, 2],
            [6, 320, 1, 1],
        ]

        # building first layer
        assert input_size % 32 == 0
        input_channel = int(32 * width_mult)
        self.last_channel = int(1280 * width_mult) if width_mult > 1.0 else 1280
        self.features = [conv_bn(3, input_channel, 2)]
        # building inverted residual blocks
        for t, c, n, s in self.inverted_residual_setting:
            output_channel = int(c * width_mult)
            for i in range(n):
                if i == 0:
                    self.features.append(InvertedResidual(input_channel, output_channel, s, t))
                else:
                    self.features.append(InvertedResidual(input_channel, output_channel, 1, t))
                input_channel = output_channel
        # building last several layers
        self.features.append(conv_1x1_bn(input_channel, self.last_channel))
        self.features.append(nn.AvgPool2d(input_size//32))
        # make it nn.Sequential
        self.features = nn.Sequential(*self.features)

        # building classifier
        self.classifier = nn.Sequential(
            nn.Dropout(),
            nn.Linear(self.last_channel, n_class),
        )

        self._initialize_weights()

    def forward(self, x):
        x = self.features(x)
        x = x.view(-1, self.last_channel)
        x = self.classifier(x)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
                if m.bias is not None:
                    m.bias.data.uniform_() * 10
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.uniform_() * 10
                m.bias.data.uniform_() * 10
            elif isinstance(m, nn.Linear):
                m.weight.data.normal_(0, 0.01)
                m.bias.data.zero_()


def save_float_tensor(float_tensor_np, save_file_path):
    assert isinstance(float_tensor_np, np.ndarray)
    assert float_tensor_np.ndim == 3

    float_tensor_pb = float_tensor_pb2.FloatTensor()
    float_tensor_pb.shape.channel = float_tensor_np.shape[0]
    float_tensor_pb.shape.height = float_tensor_np.shape[1]
    float_tensor_pb.shape.width = float_tensor_np.shape[2]
    float_tensor_pb.data.extend(float_tensor_np.flatten())

    with open(save_file_path, 'wb') as save_file:
        save_file.write(float_tensor_pb.SerializeToString())


def main():
    model = MobileNetV2()
    model.eval()

    input_x = Variable(torch.randn(1, 3, 224, 224) * 10 + 1)
    output_y = model(input_x)
    output_y = torch.unsqueeze(output_y, dim=0)
    output_y = torch.unsqueeze(output_y, dim=0)

    print('output 10 front values: ')
    print(output_y.data.numpy()[0][0, 0, 0:10])

    input_file_path = '../../test_data/mobilenet_v2_test_case_input.pb'
    output_file_path = '../../test_data/mobilenet_v2_test_case_output.pb'
    onnx_file_path = '../../test_data/mobilenet_v2_test_case.onnx'
    save_float_tensor(input_x.data.numpy()[0], input_file_path)
    save_float_tensor(output_y.data.numpy()[0], output_file_path)
    torch.onnx.export(model, input_x, onnx_file_path, verbose=False)

    subprocess.call(['../onnx2ncnn',
                     '../../test_data/mobilenet_v2_test_case.onnx',
                     '../../test_data/mobilenet_v2_test_case.param',
                     '../../test_data/mobilenet_v2_test_case.bin'])
    subprocess.call(['rm', '../../test_data/mobilenet_v2_test_case.onnx'])


if __name__ == "__main__":
    main()
